import torch
import torchvision
import numpy as np
import random
from torch import nn
from tqdm import trange, tqdm
from warmup_scheduler import GradualWarmupScheduler
import argparse
from sklearn.metrics import f1_score,average_precision_score,recall_score
import wandb
import copy

from preprocess.kfolder import kfolder_resnet_single,kfolder_resnet_double
from preprocess.traindataset import TrainDataset
from preprocess.testdataset import TestDataset
from preprocess import traindataset_double,traindataset_3d,testdataset_double,testdataset_3d
from model.get_model import get_single_model,get_double_model,get_3d_model

from sklearn.metrics import roc_auc_score,accuracy_score,multilabel_confusion_matrix




if __name__ == '__main__':

    num_classes = 5
    window_size = 640  
    batch_size = 6
    num_folder= 5 
    seed = 2022
    modelname = 'resnet50'
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print('运行设备为',device)
    train_image_path_lists,test_image_path_lists =  kfolder_resnet_single(num_folder=num_folder,image_dir='../1.数据集预处理/dataset',seed=seed,need_val=False)


    predictions = []
    labels = []
    all_pic_list = []
    for test_image_path_list in  test_image_path_lists:
        all_pic_list.extend(test_image_path_list)

    for folder_index in range(num_folder):

        pth_file = f'./pth/{modelname}_{folder_index}.pth'
        model = None

        model = get_single_model(modelname,num_classes=num_classes)

        model.load_state_dict(torch.load(pth_file,map_location=device))
        model = model.to(device)
        testdataset = TestDataset(test_image_path_lists[folder_index])
        testloader = torch.utils.data.DataLoader(testdataset, batch_size=batch_size, shuffle=False,num_workers=12)
        

        
        
        model.eval()


        for index, data in tqdm(enumerate(testloader)):
            test_img,test_label = data

            test_labels = test_label.cpu().detach().tolist()
            labels.extend(test_labels)
            
            test_inputs = test_img.to(device)
            test_outputs = model(test_inputs)
            result = np.argmax(nn.functional.softmax(test_outputs, dim=1).cpu().detach().numpy(),axis=1)
            predictions.extend(result)

        wrong_index_list = [i for i in range(len(labels)) if labels[i]!=predictions[i] ]

        print(f"共计{len(wrong_index_list)}个错误")
        
    for  wrong_index in wrong_index_list:
        print(f'错误图片名{all_pic_list[wrong_index]}，将 {labels[wrong_index]}识别为 {predictions[wrong_index]}')